import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import matplotlib

def create_architecture_diagram():
    # 创建画布
    fig, ax = plt.subplots(figsize=(15, 10))

    # 设置背景色和去除坐标轴
    ax.set_facecolor('#f5f5f5')
    ax.axis('off')

    # 定义模块颜色
    backbone_color = '#bdd7ee'  # 浅蓝
    lam_color = '#ffe699'  # 浅黄
    lcm_color = '#c6e0b4'  # 浅绿
    adem_color = '#f8cbad'  # 浅橙
    head_color = '#d9d2e9'  # 浅紫

    # 绘制模块框
    # 主干网络
    backbone = patches.Rectangle((0.2, 0.7), 0.15, 0.2,
                                 facecolor=backbone_color, edgecolor='black',
                                 linewidth=2, alpha=0.8)
    ax.add_patch(backbone)
    ax.text(0.275, 0.8, 'Backbone\nResNet-50', ha='center', va='center', fontsize=12, fontweight='bold')

    # LAM模块
    lam = patches.Rectangle((0.45, 0.7), 0.15, 0.2,
                            facecolor=lam_color, edgecolor='black',
                            linewidth=2, alpha=0.8)
    ax.add_patch(lam)
    ax.text(0.525, 0.8, 'LAM\nLocal Affinity\nModule', ha='center', va='center', fontsize=12, fontweight='bold')

    # LCM模块
    lcm = patches.Rectangle((0.7, 0.7), 0.15, 0.2,
                            facecolor=lcm_color, edgecolor='black',
                            linewidth=2, alpha=0.8)
    ax.add_patch(lcm)
    ax.text(0.775, 0.8, 'LCM\nLocal Collaborating\nModule', ha='center', va='center', fontsize=12, fontweight='bold')

    # ADEM模块
    adem = patches.Rectangle((0.45, 0.4), 0.15, 0.2,
                             facecolor=adem_color, edgecolor='black',
                             linewidth=2, alpha=0.8)
    ax.add_patch(adem)
    ax.text(0.525, 0.5, 'ADEM\nAuxiliary Density\nEstimation', ha='center', va='center', fontsize=12, fontweight='bold')

    # 检测头
    det_head = patches.Rectangle((0.7, 0.4), 0.15, 0.2,
                                 facecolor=head_color, edgecolor='black',
                                 linewidth=2, alpha=0.8)
    ax.add_patch(det_head)
    ax.text(0.775, 0.5, 'Detection\nHead', ha='center', va='center', fontsize=12, fontweight='bold')

    # 绘制连接箭头
    # Backbone -> LAM
    ax.arrow(0.35, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # LAM -> LCM
    ax.arrow(0.6, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # LCM -> Detection Head
    ax.arrow(0.775, 0.7, 0, -0.09, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # LAM -> ADEM (训练时)
    ax.arrow(0.525, 0.7, 0, -0.09, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2, linestyle='--')

    # ADEM -> Detection Head (训练时辅助监督)
    ax.arrow(0.6, 0.5, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2, linestyle='--')

    # 添加图例
    solid_line = plt.Line2D([0], [0], color='black', linewidth=2, linestyle='-')
    dashed_line = plt.Line2D([0], [0], color='black', linewidth=2, linestyle='--')
    ax.legend([solid_line, dashed_line],
              ['Inference & Training Path', 'Training Path Only'],
              loc='upper center', bbox_to_anchor=(0.5, 0.1))

    # 添加标题
    ax.set_title('Density-Aware Improved DiffusionDet Architecture', fontsize=16, fontweight='bold', pad=20)

    plt.tight_layout()
    plt.savefig('architecture_diagram.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig





def create_lam_diagram():
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_facecolor('#f9f9f9')
    ax.axis('off')

    # 定义颜色
    input_color = '#bdd7ee'
    process_color = '#ffe699'
    output_color = '#c6e0b4'

    # 绘制模块框
    # 输入特征
    input_feat = patches.Rectangle((0.1, 0.7), 0.15, 0.2,
                                   facecolor=input_color, edgecolor='black',
                                   linewidth=2, alpha=0.8)
    ax.add_patch(input_feat)
    ax.text(0.175, 0.8, 'Input\nFeature Map\nF', ha='center', va='center', fontsize=12)

    # 通道注意力
    channel_attn = patches.Rectangle((0.35, 0.8), 0.2, 0.1,
                                     facecolor=process_color, edgecolor='black',
                                     linewidth=2, alpha=0.8)
    ax.add_patch(channel_attn)
    ax.text(0.45, 0.85, 'Channel Attention', ha='center', va='center', fontsize=10)

    # 空间注意力
    spatial_attn = patches.Rectangle((0.35, 0.6), 0.2, 0.1,
                                     facecolor=process_color, edgecolor='black',
                                     linewidth=2, alpha=0.8)
    ax.add_patch(spatial_attn)
    ax.text(0.45, 0.65, 'Spatial Attention', ha='center', va='center', fontsize=10)

    # 注意力融合
    attn_fusion = patches.Rectangle((0.65, 0.7), 0.15, 0.2,
                                    facecolor=process_color, edgecolor='black',
                                    linewidth=2, alpha=0.8)
    ax.add_patch(attn_fusion)
    ax.text(0.725, 0.8, 'Attention\nFusion', ha='center', va='center', fontsize=12)

    # 输出特征
    output_feat = patches.Rectangle((0.9, 0.7), 0.15, 0.2,
                                    facecolor=output_color, edgecolor='black',
                                    linewidth=2, alpha=0.8)
    ax.add_patch(output_feat)
    ax.text(0.975, 0.8, 'Enhanced\nFeature Map\nF\'', ha='center', va='center', fontsize=12)

    # 绘制连接箭头
    ax.arrow(0.25, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)
    ax.arrow(0.25, 0.8, 0.09, -0.15, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    ax.arrow(0.55, 0.85, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)
    ax.arrow(0.55, 0.65, 0.09, 0.05, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    ax.arrow(0.8, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # 添加LAM公式
    formula = r"$F' = F + \gamma \cdot (F \otimes Attention)$"
    ax.text(0.5, 0.4, formula, ha='center', va='center', fontsize=16)

    # 添加模块说明
    ax.text(0.5, 0.3, "LAM通过结合通道注意力和空间注意力机制，\n增强局部区域内特征的一致性，有助于抑制噪声干扰",
            ha='center', va='center', fontsize=12)

    # 添加标题
    ax.set_title('Local Affinity Module (LAM) Architecture', fontsize=16, fontweight='bold', pad=20)

    plt.tight_layout()
    plt.savefig('lam_diagram.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig





def create_lcm_diagram():
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_facecolor('#f9f9f9')
    ax.axis('off')

    # 定义颜色
    input_color = '#bdd7ee'
    process_color = '#c6e0b4'
    output_color = '#f8cbad'

    # 绘制模块框
    # 输入特征
    input_feat = patches.Rectangle((0.1, 0.7), 0.15, 0.2,
                                   facecolor=input_color, edgecolor='black',
                                   linewidth=2, alpha=0.8)
    ax.add_patch(input_feat)
    ax.text(0.175, 0.8, 'ROI\nFeatures', ha='center', va='center', fontsize=12)

    # 特征降维
    feat_reducer = patches.Rectangle((0.35, 0.7), 0.15, 0.2,
                                     facecolor=process_color, edgecolor='black',
                                     linewidth=2, alpha=0.8)
    ax.add_patch(feat_reducer)
    ax.text(0.425, 0.8, 'Feature\nReduction', ha='center', va='center', fontsize=12)

    # 自注意力
    self_attn = patches.Rectangle((0.6, 0.7), 0.15, 0.2,
                                  facecolor=process_color, edgecolor='black',
                                  linewidth=2, alpha=0.8)
    ax.add_patch(self_attn)
    ax.text(0.675, 0.8, 'Self\nAttention', ha='center', va='center', fontsize=12)

    # 跨实例交互
    cross_inst = patches.Rectangle((0.85, 0.7), 0.15, 0.2,
                                   facecolor=process_color, edgecolor='black',
                                   linewidth=2, alpha=0.8)
    ax.add_patch(cross_inst)
    ax.text(0.925, 0.8, 'Cross-Instance\nInteraction', ha='center', va='center', fontsize=11)

    # 输出特征
    output_feat = patches.Rectangle((0.6, 0.4), 0.15, 0.2,
                                    facecolor=output_color, edgecolor='black',
                                    linewidth=2, alpha=0.8)
    ax.add_patch(output_feat)
    ax.text(0.675, 0.5, 'Enhanced\nCollaborative\nFeatures', ha='center', va='center', fontsize=11)

    # 绘制连接箭头
    ax.arrow(0.25, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    ax.arrow(0.5, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    ax.arrow(0.75, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    ax.arrow(0.925, 0.7, 0, -0.15, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    ax.arrow(0.925, 0.55, -0.16, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # 添加残差连接
    ax.arrow(0.175, 0.7, 0, -0.15, head_width=0.01, head_length=0.01,
             fc='black', ec='black', linewidth=1.5, linestyle='-')
    ax.arrow(0.175, 0.55, 0.485, 0, head_width=0.01, head_length=0.01,
             fc='black', ec='black', linewidth=1.5, linestyle='-')
    ax.text(0.4, 0.55, 'Residual Connection', ha='center', va='center', fontsize=10, color='gray')

    # 添加LCM公式
    formula = r"$F_{out} = F_{in} + LCM(F_{in})$"
    ax.text(0.5, 0.25, formula, ha='center', va='center', fontsize=16)

    # 添加模块说明
    ax.text(0.5, 0.15, "LCM通过残差自注意力机制促进特征间协作，\n增强局部区域内不同目标之间的信息交流",
            ha='center', va='center', fontsize=12)

    # 添加标题
    ax.set_title('Local Collaborating Module (LCM) Architecture', fontsize=16, fontweight='bold', pad=20)

    plt.tight_layout()
    plt.savefig('lcm_diagram.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig




def create_adem_diagram():
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_facecolor('#f9f9f9')
    ax.axis('off')

    # 定义颜色
    input_color = '#bdd7ee'
    process_color = '#f8cbad'
    output_color = '#d9d2e9'

    # 绘制模块框
    # 多尺度FPN特征
    fpn_feats = []
    fpn_labels = ['P3', 'P4', 'P5']
    for i in range(3):
        y_pos = 0.8 - i * 0.15
        fpn_feat = patches.Rectangle((0.1, y_pos), 0.15, 0.1,
                                     facecolor=input_color, edgecolor='black',
                                     linewidth=2, alpha=0.8)
        ax.add_patch(fpn_feat)
        ax.text(0.175, y_pos + 0.05, f'FPN {fpn_labels[i]}', ha='center', va='center', fontsize=10)
        fpn_feats.append(fpn_feat)

    # 特征转换
    feat_transforms = []
    for i in range(3):
        y_pos = 0.8 - i * 0.15
        transform = patches.Rectangle((0.35, y_pos), 0.15, 0.1,
                                      facecolor=process_color, edgecolor='black',
                                      linewidth=2, alpha=0.8)
        ax.add_patch(transform)
        ax.text(0.425, y_pos + 0.05, f'Scale\nFusion {i + 1}', ha='center', va='center', fontsize=10)
        feat_transforms.append(transform)

        # 连接箭头
        ax.arrow(0.25, y_pos + 0.05, 0.09, 0, head_width=0.02, head_length=0.02,
                 fc='black', ec='black', linewidth=2)

    # 特征融合
    feat_fusion = patches.Rectangle((0.6, 0.7), 0.15, 0.2,
                                    facecolor=process_color, edgecolor='black',
                                    linewidth=2, alpha=0.8)
    ax.add_patch(feat_fusion)
    ax.text(0.675, 0.8, 'Multi-scale\nFeature\nFusion', ha='center', va='center', fontsize=12)

    # 上采样箭头
    for i in range(3):
        y_pos = 0.8 - i * 0.15
        ax.arrow(0.5, y_pos + 0.05, 0.09, 0.75 - y_pos, head_width=0.02, head_length=0.02,
                 fc='black', ec='black', linewidth=2)

    # 密度预测
    density_pred = patches.Rectangle((0.85, 0.8), 0.15, 0.1,
                                     facecolor=output_color, edgecolor='black',
                                     linewidth=2, alpha=0.8)
    ax.add_patch(density_pred)
    ax.text(0.925, 0.85, 'Density Map', ha='center', va='center', fontsize=12)

    # 不确定性预测
    sigma_pred = patches.Rectangle((0.85, 0.65), 0.15, 0.1,
                                   facecolor=output_color, edgecolor='black',
                                   linewidth=2, alpha=0.8)
    ax.add_patch(sigma_pred)
    ax.text(0.925, 0.7, 'Uncertainty\nσ', ha='center', va='center', fontsize=12)

    # 密度预测箭头
    ax.arrow(0.75, 0.8, 0.09, 0, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # 不确定性预测箭头
    ax.arrow(0.75, 0.75, 0.09, -0.05, head_width=0.02, head_length=0.02,
             fc='black', ec='black', linewidth=2)

    # 添加密度损失公式
    formula = r"$\mathcal{L}_{density} = \frac{(D_{pred} - D_{gt})^2}{2\sigma^2} + \log \sigma + \lambda_{adem} \mathcal{L}_{reg}$"
    ax.text(0.5, 0.4, formula, ha='center', va='center', fontsize=16)

    # 添加模块说明
    ax.text(0.5, 0.3, "ADEM generates scene density maps through multi-scale feature fusion,\n providing density prior information to assist detection and adapt to dense object distributions",
            ha='center', va='center', fontsize=12)

    # 添加可视化图例
    example_density = plt.imread("/media/ross/8TB/project/lsh/dataset/microAlgea/mmdetection_format/test/0e6f88b5dca4cdc699902164229b330b5af255888811b273354bcb46b0d1d857_jpeg.rf.f52e644be91d2eea6f3da06d99f8119d.jpg") if "path/to/example_density.png" else np.random.rand(
        100, 100)
    density_inset = fig.add_axes([0.65, 0.2, 0.3, 0.1])
    density_inset.imshow(example_density, cmap='hot')
    density_inset.set_title("Density Map Example", fontsize=10)
    density_inset.axis('off')

    # 添加标题
    ax.set_title('Auxiliary Density Estimation Module (ADEM) Architecture', fontsize=16, fontweight='bold', pad=20)

    plt.tight_layout()
    plt.savefig('adem_diagram.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig




def create_performance_comparison():
    # 假设的性能数据，请替换为您的真实实验结果
    methods = ['Baseline\nDiffusionDet', '+ LAM', '+ LCM', '+ ADEM', '全部改进\n(Ours)']
    ap_values = [31.2, 32.8, 33.5, 33.9, 35.7]  # AP
    ap_small = [13.5, 16.7, 17.2, 19.8, 22.5]  # AP_small
    ap_crowded = [25.7, 28.3, 29.1, 32.2, 34.6]  # AP在密集区域

    x = np.arange(len(methods))
    width = 0.25

    fig, ax = plt.subplots(figsize=(12, 8))
    rects1 = ax.bar(x - width, ap_values, width, label='AP', color='#5b9bd5')
    rects2 = ax.bar(x, ap_small, width, label='AP$_{small}$', color='#ed7d31')
    rects3 = ax.bar(x + width, ap_crowded, width, label='AP$_{dense}$', color='#70ad47')

    # 添加数值标签
    def add_labels(rects):
        for rect in rects:
            height = rect.get_height()
            ax.annotate(f'{height}',
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),
                        textcoords="offset points",
                        ha='center', va='bottom', fontsize=10)

    add_labels(rects1)
    add_labels(rects2)
    add_labels(rects3)

    ax.set_ylabel('Accuracy (%)', fontsize=14)
    ax.set_title('Performance Comparison of Different Modules', fontsize=16, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(methods, fontsize=12)
    ax.legend(fontsize=12)
    ax.grid(axis='y', linestyle='--', alpha=0.7)

    # 增加改进百分比标签
    for i in range(1, len(methods)):
        improvement = ((ap_values[i] - ap_values[0]) / ap_values[0]) * 100
        ax.annotate(f'+{improvement:.1f}%',
                    xy=(i - width, ap_values[i] + 1),
                    color='blue',
                    fontsize=9)

    plt.tight_layout()
    plt.savefig('performance_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig




def create_density_curve():
    # 假设的不同密度下的性能数据
    density_levels = [10, 20, 30, 40, 50, 60, 70, 80]  # 目标密度（目标数/图像）
    baseline_ap = [40.5, 38.2, 33.7, 28.5, 24.3, 21.8, 19.5, 17.2]  # 基线模型AP
    our_ap = [41.0, 39.5, 36.8, 33.2, 30.5, 28.7, 26.9, 25.1]  # 我们改进的模型AP

    fig, ax = plt.subplots(figsize=(10, 7))

    ax.plot(density_levels, baseline_ap, 'o-', linewidth=2, markersize=8, label='DiffusionDet', color='#5b9bd5')
    ax.plot(density_levels, our_ap, 's-', linewidth=2, markersize=8, label='Ours (with LAM+LCM+ADEM)', color='#70ad47')

    # 填充两条曲线之间的区域以突出性能提升
    ax.fill_between(density_levels, baseline_ap, our_ap, color='#ddebf7', alpha=0.5)

    # 标注关键点的性能差异
    for i in [2, 4, 6]:
        diff = our_ap[i] - baseline_ap[i]
        ax.annotate(f'+{diff:.1f}',
                    xy=(density_levels[i], (baseline_ap[i] + our_ap[i]) / 2),
                    xytext=(5, 0),
                    textcoords="offset points",
                    fontsize=10,
                    bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))

    ax.set_xlabel('Scene Density (Objects/Image)', fontsize=14)
    ax.set_ylabel('Average Precision AP (%)', fontsize=14)
    ax.set_title('Detection Performance under Different Scene Densities', fontsize=16, fontweight='bold')
    ax.grid(True, linestyle='--', alpha=0.7)
    ax.legend(fontsize=12)

    # 添加密集区域标记
    ax.axvspan(50, 80, alpha=0.15, color='red')
    ax.text(65, 42, 'Dense Region\n(High Challenge)', fontsize=12, ha='center',
            bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="red", alpha=0.7))

    plt.tight_layout()
    plt.savefig('density_curve.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig


def create_detection_visualization():
    # 这个函数需要真实图像和检测结果
    # 以下是模拟代码，您需要替换为实际图像和检测结果

    import matplotlib.patches as patches
    from matplotlib.colors import LinearSegmentedColormap

    fig, axes = plt.subplots(2, 2, figsize=(15, 12))

    # 假设有两个测试图像
    for row in range(2):
        # 加载图像(替换为您的真实图像路径)
        if row == 0:
            img_path = "/media/ross/8TB/project/lsh/dataset/microAlgea/mmdetection_format/test/0e6f88b5dca4cdc699902164229b330b5af255888811b273354bcb46b0d1d857_jpeg.rf.f52e644be91d2eea6f3da06d99f8119d.jpg"  # 密集目标场景
        else:
            img_path = "/media/ross/8TB/project/lsh/dataset/microAlgea/mmdetection_format/test/43_14_jpg.rf.33dcb9e2bc20906ffce4d14f5f0c560b.jpg"  # 另一密集目标场景

        # 为了演示，创建随机图像
        img = np.random.rand(300, 400, 3)

        # 第一列显示基线方法结果
        axes[row, 0].imshow(img)
        axes[row, 0].set_title(f"{'Dense Scene' if row == 0 else 'Overlapping Objects'} - Baseline DiffusionDet", fontsize=12)
        axes[row, 0].axis('off')

        # 添加随机基线方法检测框(红色)
        np.random.seed(42 + row)  # 固定随机种子以确保可重复性
        n_boxes = 15 if row == 0 else 10
        for i in range(n_boxes):
            x, y = np.random.randint(50, 300), np.random.randint(50, 200)
            w, h = np.random.randint(30, 80), np.random.randint(30, 80)
            rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='red', facecolor='none')
            axes[row, 0].add_patch(rect)
            # 随机添加一些错误或漏检的框
            if i % 5 == 0:
                axes[row, 0].text(x, y - 5, "Missed", color='red', fontsize=10,
                                  bbox=dict(facecolor='white', alpha=0.7))

        # 第二列显示我们的改进方法结果
        axes[row, 1].imshow(img)
        axes[row, 1].set_title(f"{'Dense Scene' if row == 0 else 'Overlapping Objects'} - Our Method(LAM+LCM+ADEM)", fontsize=12)
        axes[row, 1].axis('off')

        # 添加改进方法的检测框(绿色)并可视化密度
        np.random.seed(42 + row)  # 使用相同的随机种子以确保可比性

        # 添加密度热力图可视化
        density = np.zeros((img.shape[0], img.shape[1]))
        for i in range(n_boxes + 5):  # 我们的方法检出更多目标
            x, y = np.random.randint(50, 300), np.random.randint(50, 200)
            w, h = np.random.randint(30, 80), np.random.randint(30, 80)

            # 在目标周围添加高斯密度
            y_grid, x_grid = np.mgrid[0:img.shape[0], 0:img.shape[1]]
            sigma = 30
            density += np.exp(-((x_grid - x - w / 2) ** 2 + (y_grid - y - h / 2) ** 2) / (2 * sigma ** 2))

            rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='green', facecolor='none')
            axes[row, 1].add_patch(rect)

        # 归一化密度图
        density = density / density.max()

        # 创建透明度蒙版，只在密度较高的区域显示热力图
        alpha_mask = np.zeros_like(density)
        alpha_mask[density > 0.3] = 0.4  # 设置透明度

        # 创建自定义颜色映射
        colors = [(0, 0, 0, 0), (0, 0, 1, 0.3), (0, 1, 1, 0.3), (0, 1, 0, 0.3), (1, 1, 0, 0.3), (1, 0, 0, 0.3)]
        cmap = LinearSegmentedColormap.from_list('density_cmap', colors, N=256)

        # 绘制密度热力图
        axes[row, 1].imshow(density, cmap=cmap, alpha=alpha_mask, vmin=0, vmax=1)

    # 添加总体标题和说明
    fig.suptitle('Detection Results Visualization Comparison', fontsize=18, fontweight='bold', y=0.98)
    fig.text(0.5, 0.01, 'Density-Aware Improved DiffusionDet can more accurately detect objects in dense scenes, reducing missed detections and false positives',
             ha='center', fontsize=14)

    plt.tight_layout()
    plt.subplots_adjust(top=0.93, bottom=0.05)
    plt.savefig('detection_visualization.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig


# 注意：真实使用时需替换图像路径和实际的检测结果
# create_detection_visualization()


def create_error_analysis():
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))

    # 错误类型标题
    error_types = ['Small Object Misses', 'Dense Region False Positives', 'Overlapping Object Confusion']

    # 为了演示，创建随机图像
    for row in range(2):
        for col in range(3):
            # 随机图像
            img = np.random.rand(200, 200, 3)
            axes[row, col].imshow(img)

            if row == 0:
                # 第一行显示基线方法的错误
                method = "Baseline Method"
                box_color = 'red'
            else:
                # 第二行显示我们方法的改进
                method = "Our Method"
                box_color = 'green'

            axes[row, col].set_title(f"{error_types[col]}\n({method})", fontsize=12)
            axes[row, col].axis('off')

            # 根据错误类型添加不同的可视化
            if col == 0:  # 小目标漏检
                # 添加几个小目标
                for i in range(5):
                    x, y = np.random.randint(50, 150), np.random.randint(50, 150)
                    w, h = np.random.randint(10, 20), np.random.randint(10, 20)
                    rect = patches.Rectangle((x, y), w, h, linewidth=2,
                                             edgecolor=box_color if row == 1 or i < 2 else 'none',
                                             facecolor='none')
                    axes[row, col].add_patch(rect)
                    if row == 0 and i >= 2:
                        axes[row, col].text(x, y - 5, "Missed", color='red', fontsize=10)

            elif col == 1:  # 密集区域误检
                # 添加密集区域
                for i in range(10):
                    x, y = 70 + (i % 3) * 20, 70 + (i // 3) * 20
                    w, h = np.random.randint(15, 25), np.random.randint(15, 25)
                    rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor=box_color, facecolor='none')
                    axes[row, col].add_patch(rect)

                # 在基线方法中添加误检
                if row == 0:
                    for i in range(2):
                        x, y = 85 + i * 30, 85 + i * 30
                        w, h = np.random.randint(15, 25), np.random.randint(15, 25)
                        rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='red', facecolor='none',
                                                 linestyle='--')
                        axes[row, col].add_patch(rect)
                        axes[row, col].text(x, y - 5, "False Positive", color='red', fontsize=10)

            else:  # 重叠目标混淆
                # 添加重叠目标
                x1, y1 = 70, 70
                w1, h1 = 70, 70
                x2, y2 = 100, 100
                w2, h2 = 70, 70

                # 第一个目标
                rect1 = patches.Rectangle((x1, y1), w1, h1, linewidth=2, edgecolor=box_color, facecolor='none')
                axes[row, col].add_patch(rect1)

                # 第二个目标
                rect2 = patches.Rectangle((x2, y2), w2, h2, linewidth=2, edgecolor=box_color, facecolor='none')
                axes[row, col].add_patch(rect2)

                # 基线方法中的混淆
                if row == 0:
                    rect_confused = patches.Rectangle((x1 + 20, y1 + 20), w1 + 30, h1 + 30, linewidth=2,
                                                      edgecolor='red',
                                                      facecolor='none', linestyle='--')
                    axes[row, col].add_patch(rect_confused)
                    axes[row, col].text(x1 + 25, y1 + 15, "Confused", color='red', fontsize=10)

    # 添加总体标题和错误分析说明
    fig.suptitle('Typical Error Analysis and Improvement', fontsize=16, fontweight='bold', y=0.98)
    fig.text(0.5, 0.01, 'Our density-aware method significantly improves small object detection, dense region accuracy, and overlapping object discrimination',
             ha='center', fontsize=14)

    plt.tight_layout()
    plt.subplots_adjust(top=0.92, bottom=0.08)
    plt.savefig('error_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig


def create_diffusion_process_visualization():
    fig, axes = plt.subplots(2, 4, figsize=(16, 8))

    # 设置图表标题
    timesteps = ["Initial Noise", "t=800", "t=400", "t=0 (Final Result)"]

    # 为了演示，创建随机图像
    img = np.random.rand(200, 200, 3)

    for row in range(2):
        for col in range(4):
            # 在每个子图中显示图像
            axes[row, col].imshow(img)
            axes[row, col].axis('off')

            if row == 0:
                # 第一行展示原始DiffusionDet的扩散过程
                subtitle = f"DiffusionDet: {timesteps[col]}"
            else:
                # 第二行展示我们改进的扩散过程
                subtitle = f"Improved DiffusionDet: {timesteps[col]}"

            axes[row, col].set_title(subtitle, fontsize=12)

            # 根据时间步长添加不同的框
            n_boxes = 50 if col == 0 else int(50 * (1 - col / 4))

            for i in range(n_boxes):
                x = np.random.randint(0, 180)
                y = np.random.randint(0, 180)
                w = np.random.randint(10, 50)
                h = np.random.randint(10, 50)

                # 根据时间步长调整框的清晰度
                alpha = 0.2 if col == 0 else 0.3 + 0.7 * (col / 3)
                linewidth = 1 if col == 0 else 1 + col

                # 为我们的方法添加密度感知优势
                if row == 1:
                    if col >= 2:  # 在后期时间步长中表现出优势
                        # 调整框的位置以表示更精确的定位
                        if i % 5 == 0:
                            x = x + np.random.randint(-10, 10)
                            y = y + np.random.randint(-10, 10)

                        # 为密集区域添加更多框
                        if x > 100 and y > 100 and i % 3 == 0 and col == 3:
                            x2 = x + np.random.randint(-20, 20)
                            y2 = y + np.random.randint(-20, 20)
                            rect2 = patches.Rectangle((x2, y2), w, h, linewidth=linewidth,
                                                      edgecolor='green', facecolor='none',
                                                      alpha=alpha, linestyle='-')
                            axes[row, col].add_patch(rect2)

                # 添加扩散过程中的边界框
                if col < 3:  # 非最终结果时显示为虚线
                    linestyle = '--'
                    color = 'blue' if row == 0 else 'green'
                else:  # 最终结果显示为实线
                    linestyle = '-'
                    color = 'red' if row == 0 else 'green'

                rect = patches.Rectangle((x, y), w, h, linewidth=linewidth,
                                         edgecolor=color, facecolor='none',
                                         alpha=alpha, linestyle=linestyle)
                axes[row, col].add_patch(rect)

    # 如果有时间，可以添加密度图在后期时间步长
    if True:  # 始终执行
        # 创建简单的密度图
        density = np.zeros((200, 200))
        for i in range(30):
            x, y = np.random.randint(50, 150), np.random.randint(50, 150)
            sigma = 20
            y_grid, x_grid = np.mgrid[0:200, 0:200]
            density += np.exp(-((x_grid - x) ** 2 + (y_grid - y) ** 2) / (2 * sigma ** 2))

        # 归一化密度图
        density = density / density.max()

        # 只在我们方法的后期步骤中显示密度图
        for col in [2, 3]:  # t=400和t=0
            cmap = plt.cm.hot
            axes[1, col].imshow(density, cmap=cmap, alpha=0.3)

    # 添加总体标题和说明
    fig.suptitle('Detection Box Optimization During Diffusion Process', fontsize=16, fontweight='bold')
    fig.text(0.5, 0.02, 'Improved DiffusionDet utilizes density-aware mechanisms to more accurately optimize object positions and boundaries during the diffusion process',
             ha='center', fontsize=14)

    plt.tight_layout()
    plt.subplots_adjust(top=0.9, bottom=0.1)
    plt.savefig('diffusion_process.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig


def create_network_architecture():
    fig, ax = plt.subplots(figsize=(15, 10))
    ax.axis('off')

    # 创建一个空白画布
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 60)

    # 定义组件样式
    box_style = dict(boxstyle="round,pad=0.5", fc="white", ec="black", lw=2)
    improved_style = dict(boxstyle="round,pad=0.5", fc="#e6f2ff", ec="#3366cc", lw=2)
    backbone_style = dict(boxstyle="round,pad=0.5", fc="#f2f2f2", ec="black", lw=2)

    # 绘制主干网络
    ax.text(10, 55, "Input Image", fontsize=12, ha="center", va="center", bbox=box_style)
    ax.text(10, 45, "Backbone\n(ResNet/Swin)", fontsize=12, ha="center", va="center", bbox=backbone_style)

    # 原始DiffusionDet组件
    ax.text(10, 30, "Feature Extraction", fontsize=12, ha="center", va="center", bbox=box_style)
    ax.text(10, 15, "Original DiffusionDet Detection Head", fontsize=12, ha="center", va="center", bbox=box_style)
    ax.text(10, 5, "Detection Output", fontsize=12, ha="center", va="center", bbox=box_style)

    # 连接线 - 反向箭头 (xytext和xy位置互换)
    ax.annotate("", xy=(10, 48), xytext=(10, 52), arrowprops=dict(arrowstyle="->", lw=1.5))
    ax.annotate("", xy=(10, 33), xytext=(10, 42), arrowprops=dict(arrowstyle="->", lw=1.5))
    ax.annotate("", xy=(10, 18), xytext=(10, 27), arrowprops=dict(arrowstyle="->", lw=1.5))
    ax.annotate("", xy=(10, 8), xytext=(10, 12), arrowprops=dict(arrowstyle="->", lw=1.5))

    # 我们的改进组件 - LAM
    ax.text(40, 45, "LAM\n(Local Affinity Module)", fontsize=12, ha="center", va="center",
            bbox=improved_style)

    # LAM内部结构
    lam_components = [
        (33, 37, "Channel Attention"),
        (40, 37, "Spatial Attention"),
        (47, 37, "Feature Fusion")
    ]

    for x, y, text in lam_components:
        ax.text(x, y, text, fontsize=10, ha="center", va="center",
                bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#3366cc", lw=1))

    # 添加LAM连接线 - 反向箭头
    ax.annotate("", xy=(32, 45), xytext=(20, 45), arrowprops=dict(arrowstyle="->", lw=1.5))

    # 我们的改进组件 - LCM
    ax.text(40, 30, "LCM\n(Local Collaborating Module)", fontsize=12, ha="center", va="center",
            bbox=improved_style)

    # LCM内部结构
    lcm_components = [
        (33, 22, "Feature Reduction"),
        (40, 22, "Self-Attention"),
        (47, 22, "Cross-Instance Interaction")
    ]

    for x, y, text in lcm_components:
        ax.text(x, y, text, fontsize=10, ha="center", va="center",
                bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#3366cc", lw=1))

    # 添加LCM连接线 - 反向箭头
    ax.annotate("", xy=(40, 33), xytext=(40, 42), arrowprops=dict(arrowstyle="->", lw=1.5))

    # 我们的改进组件 - ADEM
    ax.text(70, 30, "ADEM\n(Auxiliary Density Estimation Module)", fontsize=12, ha="center", va="center",
            bbox=improved_style)

    # ADEM内部结构
    adem_components = [
        (63, 22, "Density Prediction"),
        (70, 22, "Uncertainty Estimation"),
        (77, 22, "Adaptive Weighting")
    ]

    for x, y, text in adem_components:
        ax.text(x, y, text, fontsize=10, ha="center", va="center",
                bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#3366cc", lw=1))

    # 添加ADEM连接线 - 反向箭头
    ax.annotate("", xy=(65, 35), xytext=(50, 45), arrowprops=dict(arrowstyle="->", lw=1.5, ls=":"))

    # 将所有改进组件与原始DiffusionDet连接 - 反向箭头
    ax.annotate("", xy=(32, 30), xytext=(10, 25), arrowprops=dict(arrowstyle="->", lw=1.5))
    ax.annotate("", xy=(20, 18), xytext=(48, 30), arrowprops=dict(arrowstyle="->", lw=1.5))

    # 添加ADEM到损失函数的连接 - 反向箭头
    ax.text(70, 15, "Joint Optimization Loss Function", fontsize=12, ha="center", va="center",
            bbox=dict(boxstyle="round,pad=0.5", fc="#fff2e6", ec="#ff9933", lw=2))
    ax.annotate("", xy=(70, 18), xytext=(70, 27), arrowprops=dict(arrowstyle="->", lw=1.5, ls=":"))
    # 这里保留原来的箭头方向，因为已经是<-
    ax.annotate("", xy=(63, 15), xytext=(20, 15), arrowprops=dict(arrowstyle="<-", lw=1.5, ls=":"))

    # 添加说明文本
    ax.text(50, 55, "Improved DiffusionDet Architecture Diagram", fontsize=16, fontweight='bold', ha="center")
    ax.text(50, 2, "Solid line: Inference Path   Dashed line: Training Only", fontsize=12, ha="center")

    # 添加图例
    legend_elements = [
        plt.Rectangle((0, 0), 1, 1, fc="#f2f2f2", ec="black", label="Original Components"),
        plt.Rectangle((0, 0), 1, 1, fc="#e6f2ff", ec="#3366cc", label="Improved Modules"),
        plt.Rectangle((0, 0), 1, 1, fc="#fff2e6", ec="#ff9933", label="Optimization Components")
    ]
    ax.legend(handles=legend_elements, loc="lower right", framealpha=1)

    # 添加数据流指示
    flow_points = np.array([
        [10, 60], [10, 45], [10, 35], [25, 45], [40, 45], [40, 35],
        [40, 25], [15, 20], [15, 15], [15, 10], [15, 5]
    ])
    flow_curve = np.zeros((100, 2))
    flow_curve[:11] = flow_points

    plt.tight_layout()
    plt.savefig('network_architecture.png', dpi=300, bbox_inches='tight')
    plt.show()

    return fig



create_architecture_diagram()
create_lam_diagram()
create_lcm_diagram()
# 如果有示例密度图，请替换"path/to/example_density.png"
create_adem_diagram()
create_performance_comparison()
create_detection_visualization()
create_error_analysis()
create_diffusion_process_visualization()
create_network_architecture()


# 概念架构图 - 思维导图清晰展示了三大核心模块及其功能
# 性能曲线图 - 展示在不同密度场景下的性能改进
# 检测效果对比图 - 直观展示改进前后的检测效果差异
# 模型架构细节图 - 展示网络模块的详细连接关系
# 错误分析图 - 展示典型错误及改进效果
# 扩散过程可视化 - 展示检测过程中边界框的优化过程
# 模块交互图 - 展示三大核心模块的协同工作机制
# 实验设置总结图 - 总结实验配置和评估指标